#!/usr/bin/env python

import sys
import os
import subprocess
import re
import argparse
import psycopg2
import gzip
import sblab

parser = argparse.ArgumentParser(description= """
DESCRIPTION
    Query sblab with the supplied service_id (SLX number) to get the corresponding fastq files.
    Fastqfiles are renamed and demultiplxed as appropriate using the information in sblab.
    Output (demultiplexed) files are sent to sblab.main.fastqfiles and the demultiplxed report
    sent to demux_report.
    
    IMPORTANT: It is essential that get_fastq_from_slx can link the supplied service_id to the corresponding
    library(s). So make sure table lib2seq has the required service_id.
    
EXAMPLE
    get_fastq_from_slx.py SLX-5211
    get_fastq_from_slx.py -f myfastq.txt.gz SLX-5211 ## Process myfastq assigned to service SLX-5211 
    get_fastq_from_slx.py SLX-5211 -t 'test-import' ## Output files will have the 'test-import' tag in their name
    
""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('--service', '-s',
                   required= True,
                   help='''SLX number to get fastq files from.
                   ''')

parser.add_argument('--fastq', '-f',
                    type= str,
                    default= False,
                    required= False,
                    nargs= '+',
                   help='''List of fastq files already downloaded and to be converted
to Sanger, demultiplxed and archived as per information from the service_id.
This option will skip the downloading & renaming steps so make sure these files
have appropriate names. NB: Fastq files are expected to be gzipped and ending in
.gz
                   ''')

parser.add_argument('--filetag', '-t',
                    type= str,
                    default= '',
                    required= False,
                    help='''An optional tag to include in the filenames. E.g.
'test'. It is inserted as <slx-id>.<run-id>.<tag>.<filename>. It is probably best
to use the tag only for testing. Ignored if --fastq is passed.
                   ''')

parser.add_argument('--no_db_update',
                   action= 'store_true', 
                   help='''Do not update sblab tables database with fastqfile(s)
demux_report etc. (Use this option for testing only)
                   ''')

parser.add_argument('--no_fastqc',
                   action= 'store_true', 
                   help='''Do not execute fastqc on the generated fastq files.
                   ''')

args = parser.parse_args()

# --------------------------[ Some checks ]------------------------------
## Check service id exists
conn= psycopg2.connect(sblab.get_psycopgpass())
cur= conn.cursor()
sql= 'select * from lib2seq where service_id = %s limit 1'
cur.execute(sql, (args.service,))
slx= cur.fetchall()
if len(slx) == 0:
    sys.exit("""\nservice_id '%s' not found in table lib2seq.
It is necessary to have a service_id linking files to libraries.\n""" %(args.service))

## Check if I'm on the lustre file system
mypath= os.getcwd()
if mypath.startswith('/lustre/'):
    bsub= True
else:
    bsub= False

# --------------------------[ Download files ]--------------------------------
if not args.fastq:
    print('\nDownloading files for service %s...' %(args.service))
    fastq= sblab.downloadService(args.service, tag= args.filetag)
    print('\nFiles downloaded:\n' + '\n'.join(fastq))
else:
    fastq= args.fastq
    ## Check all files exist
    for f in fastq:
        test_fastq= open(f)
        test_fastq.close()
# ----------------------------[Encoding]--------------------------------------
print('\n\nDetermining encoding:')
encoding= []
for x in fastq:
    e= sblab.get_fastq_encoding(x)
    print(x + '\t' + e)
    encoding.append(e)
    
# -----------------------------------------------------------------------------
## At this point you have fastq files downloaded and the encoding determined. Files have been renamed
## from illumina s_1_1_sequence.txt.gz to SLX-0000.RUNID.s_1_1.txt.gz 
if sblab.is_multiplexed(args.service):
    """ Service is multilpexed, get a sample sheet for demultiplexing for each fastqfile."""
    print('\nPreparing for demultiplexing...')
    for fq in fastq:
        cmdfile= fq.rstrip('.gz').rstrip('.txt') + '.sh'  ## This file will execute all the required steps on front node or bsub
        shfh= open(cmdfile, 'w')
        shfh.write('#!/bin/sh' + '\n')
        shfh.write('set -e' + '\n\n')
        # ----------------------[ Prepare demux sheet ]-------------------------
        demux= sblab.get_demux_sheet(args.service, fq)
        demux_file= fq.rstrip('.gz').rstrip('.txt') + '.demux'
        fout= open(demux_file, 'w')
        report= fq.rstrip('.gz').rstrip('.txt') + '.demux_report'
        for line in demux:
            fout.write(line[0] + '\n')
        fout.close()
        shfh.write('demux_fuzzy.py --fastq %(fq)s --samplesheet %(demux_file)s --report %(report)s;  \n' %{'fq': fq, 'demux_file':demux_file, 'report':report})
        if not args.no_db_update:
            shfh.write('''python -c "import sblab; sblab.uplod_demux_fuzzy_report('%(report)s')";  \n''' %{'report':report})
        # ---------------------------------------------------------------------
        # commands to FastQC each file
        if not args.no_fastqc:
            fastqc_list= []
            n= 0
            fastqc_cmd= "\nfastqc_md5.py --fastqc ' -t 8 --noextract' -i "
            for line in demux:
                demuxfq= line[2] + '.gz'
                if n < 8: ## Number of files processed by the same fastqc run
                    fastqc_cmd= fastqc_cmd + demuxfq + ' '
                else:
                    fastqc_list.append(fastqc_cmd)
                    fastqc_cmd= "\nfastqc_md5.py --fastqc ' -t 8 --noextract' -i " + demuxfq + ' '
                    n= 0
                n += 1
            fastqc_list.append(fastqc_cmd)
            fastqc_cmd= '\n'.join(fastqc_list) + '\n\n'
            shfh.write(fastqc_cmd)
        # ---------------------------------------------------------------------
        # Commads to upload fastqc to sblab
        if not args.no_db_update and not args.no_fastqc:
            send_fastqc= []
            for line in demux:
                fastqc= sblab.get_fastqc_dir(line[2]) + '.zip'
                upcmd= 'fastqc_to_pgtable.py -i %s;' %(fastqc)
                send_fastqc.append(upcmd)
            send_fastqc= '\n'.join(send_fastqc)
            shfh.write(send_fastqc)
        # ---------------------------------------------------------------------
        # Commands to upload each demultiplexed file to sblab
        if not args.no_db_update:
            upload_fastq= []
            for line in demux:
                demuxfq= line[2]
                libid= line[1]
                upcmd= 'insert_fastq.py -i %s %s.gz;' %(libid, demuxfq)
                upload_fastq.append(upcmd)
            upload_fastq= '\n'.join(upload_fastq)
            shfh.write(upload_fastq)        
        # ---------------------------------------------------------------------
        shfh.write('\n\nexit')
        shfh.close()
        # ---------------------------------------------------------------------
        if bsub:
            cmd= 'bsub -J get_fastq_from_slx_%(fq)s -oo %(fq)s.bsub.log < %(cmdfile)s' %{'fq':fq, 'cmdfile':cmdfile}
        else:
            cmd= 'sh %s' %(cmdfile)            
        print('Executing:\n' + cmd)
        p= subprocess.Popen(cmd, shell= True)
        p.wait()
else:
    """Service is not multiplexed """
    library= sblab.get_library_for_service(args.service)
    if len(library) > 1:
        sys.exit('More than one library_id found for non-multiplexed service %s: %s!' %(args.service, library))
    library= library[0] ## Make the only library a string
    for fq, e in zip(fastq, encoding):
        if args.fastq:
            lib_fq= fq
        else:
            ## Rename only if the files have been downloaded from lims, otherwise assume the files have the correct name
            lib_fq= library + '.' + fq
            print("\nRenaming %s to %s" %(fq, lib_fq))
            os.rename(fq, lib_fq)
        cmdfile= fq.rstrip('.gz').rstrip('.txt') + '.sh'
        shfh= open(cmdfile, 'w')
        shfh.write('#!/bin/sh' + '\n')
        shfh.write('set -e' + '\n')
        if e != 'Sanger':
            print('\nConverting to Sanger encoding.')
            fastq_unzipped= lib_fq.strip('.gz')
            newname_unzipped= fastq_unzipped.strip('.txt')
            if not newname_unzipped.endswith('.fq') or not newname_unzipped.endswith('.fastq'):
                newname_unzipped= newname_unzipped + '.fq'
            newname_zipped= newname_unzipped + '.gz'
            shcmd= 'gunzip %(lib_fq)s; solexa2phred %(fastq_unzipped)s %(newname_unzipped)s; gzip %(newname_unzipped)s; rm %(fastq_unzipped)s' %{'lib_fq':lib_fq, 'fastq_unzipped':fastq_unzipped, 'newname_unzipped':newname_unzipped, 'fastq_unzipped':fastq_unzipped}
            shfh.write(shcmd + '\n')
        else:
            newname_zipped= lib_fq
        if not args.no_fastqc:
            shfh.write("\nfastqc_md5.py --fastqc ' --noextract' -i %s \n" %(newname_zipped))
        if not args.no_fastqc and not args.no_db_update:
            fastqc= sblab.get_fastqc_dir(newname_zipped) + '.zip'
            shfh.write('fastqc_to_pgtable.py -i %s \n' %(fastqc))            
        if not args.no_db_update:
            "Service is not multiplexed and file is already in sanger quality and it has been renamed. Just upload it"
            shfh.write('insert_fastq.py %s' %(newname_zipped) + '\n')
        shfh.write('exit')
        shfh.close()
        if bsub:
            cmd= 'bsub -J get_fastq_from_slx_%(fq)s -oo %(fq)s.bsub.log < %(cmdfile)s' %{'fq':fq, 'cmdfile':cmdfile}
        else:
            cmd= 'sh %s' %(cmdfile)
        print('Executing:\n' + cmd)
        p= subprocess.Popen(cmd, shell= True)
        p.wait()        
sys.exit()

